import numpy as np
import cv2
import os
import math
from datetime import datetime
from PIL import Image
import pandas as pd
import imageio
from lib import Heading


# 基本参数
start = np.array([0.0, 0.0])
goal = np.array([10.0, 10.0])
obstacles = [np.array([5.0, 5.0]), np.array([6.0, 8.0]), np.array([7.0, 5.0])]
obstacle_vel = [np.array([0.1, 0.0]), np.array([-0.05, -0.05]), np.array([0.0, 0.1])]
dt = 0.1
eta = 2.0
zeta = 10.0
d_safe = 0.5
velocity_max = 1.5

# 创建输出目录
today = datetime.now().strftime('%Y%m%d_%H%M%S')
out_dir = f"APF_Output_{today}"
os.makedirs(out_dir, exist_ok=True)

trajectory = [start.copy()]
vel_heading = []
frames = []


def attractive_force(pos, goal, alpha=10.0, beta=0.2):
    diff = goal - pos
    dist = np.linalg.norm(diff)
    if dist == 0:
        return np.zeros(2)
    direction = diff / dist
    force = alpha * (1 - np.exp(-beta * dist)) * direction
    return force

def improved_repulsive_force(pos, obstacles, l_ob=2.5, c_ob=2.0):
    F_rep = np.zeros(2)
    for obs in obstacles:
        diff = pos - obs
        dist = np.linalg.norm(diff)
        if dist == 0:
            continue
        # 抑制震荡项（你给出的权重函数）
        sigma = (1 / l_ob) * np.exp(-l_ob * dist + c_ob)
        # 原始斥力项
        if dist <= d_safe:
            F = zeta * np.exp(-dist / d_safe) * (diff / dist)
            F_rep += sigma * F
    return F_rep

def draw_scene(pos, obstacles, goal, trajectory):
    img = np.ones((600, 600, 3), dtype=np.uint8) * 255
    scale = 50
    # def trans(p): return (int(p[0] * scale), int(600 - p[1] * scale))

    def trans(p): return (int(p[0] * scale), int(p[1] * scale))

    # 起点终点
    cv2.circle(img, trans(start), 5, (255, 0, 0), -1)
    cv2.circle(img, trans(goal), 5, (0, 255, 0), -1)

    # 绘制轨迹
    for i in range(1, len(trajectory)):
        cv2.line(img, trans(trajectory[i-1]), trans(trajectory[i]), (0, 0, 255), 2)

    # 当前点
    cv2.circle(img, trans(pos), 5, (0, 0, 255), -1)

    # 障碍物
    for obs in obstacles:
        cv2.circle(img, trans(obs), int(d_safe * scale), (0, 0, 0), 2)
        cv2.circle(img, trans(obs), 3, (0, 0, 0), -1)

    return img

# 主循环
pos = start.copy()
heading_records = []
speed_records = []
while np.linalg.norm(pos - goal) > 0.5:
    F_attr = attractive_force(pos, goal)
    F_rep = improved_repulsive_force(pos, obstacles)
    F_total = F_attr + F_rep

    norm = np.linalg.norm(F_total)
    if norm > 1e-5:
        velocity = F_total / norm * velocity_max
    else:
        velocity = np.zeros(2)
    # velocity = F_total * dt
    speed = np.linalg.norm(velocity)
    heading = math.degrees(math.atan2(velocity[1], velocity[0]))
    vel_heading.append((speed, heading))
    heading_records.append(heading)
    speed_records.append(speed)
    pos = pos + velocity * dt
    trajectory.append(pos.copy())

    for i in range(len(obstacles)):
        obstacles[i] += obstacle_vel[i] * dt

    frame = draw_scene(pos, obstacles, goal, trajectory)
    frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
    cv2.imshow('APF Simulation', frame)
    if cv2.waitKey(1) & 0xFF == 27:
        break

cv2.destroyAllWindows()

# 保存动图和速度数据
gif_path = os.path.join(out_dir, 'trajectory.gif')
# frames[0].save(gif_path, save_all=True, append_images=frames[1:], duration=0.05, loop=0)
imageio.mimsave(gif_path, frames, duration=0.05)

evaluator = Heading.PathSmoothnessEvaluator(heading_records, unit='deg')
result = evaluator.evaluate()
print(f"累计转角: {result['total_turning_angle']:.2f}°")
print(f"最大单次转角: {result['max_turning_angle']:.2f}°")

timestamps = [i for i in range(len(heading_records))]
df_records = pd.DataFrame({
    'Time': timestamps,
    'Speed': speed_records,
    'Yaw': heading_records

})
df_records.to_excel(os.path.join(out_dir, 'velocity_heading.xlsx'), index=False)

print(f"✅ 动图和航速航向数据已保存至文件夹：{out_dir}")